import numpy as np
import pandas as pd
import Env.EnvMain_Attack as ENVM    # 该环境下地面目标不会阵亡
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
import multiprocessing as mp
import os
import Attack_select.MySQLFunc_11 as MSF11
import json
import sys
from itertools import combinations

os.environ["OMP_NUM_THREADS"] = "1"
Data=json.load(open('Attack_select/Configs/Config_Generator.json','r',encoding='UTF-8'))
MAX_EP,DatabaseName,Step,Span,BufferSpan,N_ammo=Data.values()

class DataGenerate_1v1(mp.Process):
    def __init__(self,n_friend,n_enemy,EP,name):
        super(DataGenerate_1v1,self).__init__()
        self.N_F=n_friend
        self.N_E=n_enemy
        self.env=ENVM.ACEnv(n_friend,n_enemy)
        self.G_EP=EP
        self.name='P'+str(name)

    def run(self):
        sql = MSF11.Mysql(self.name)
        record_buffer = []
        count = 1

        while self.G_EP.value<MAX_EP:
            E,TE,index = 0.,0,0
            strategy=self.First_step_strategy()
            att_points=list(combinations(range(0,Step,Span),2))

            for i in range(len(att_points)):
                self.env.reset()
                ob = self.env.FirstStep(strategy)  # 随机初始位置
                record=(ob[0][0][0],ob[0][0][1],ob[0][0][2],ob[1][0][0],ob[1][0][1],ob[1][0][2])

                if self.name=='P0':
                    self.env.move()
                    fig = plt.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    self.env.render(fig, ax)

                for j in range(Step):
                    ob = self.env.step(self.decision(att_points[i],j))
                    if self.name=='P0':
                        self.env.render(fig, ax)
                    E += ob[0][2][0]    # 一回合的摧毁目标期望
                if self.name=='P0':
                    plt.close()
                if E>TE:
                    TE=E
                    index=i
                E=0
                record+=(index,TE)

            record_buffer.append(record)
            if count%BufferSpan==0:
                sql.Insert_record(DatabaseName,record_buffer)
                record_buffer.clear()
            with self.G_EP.get_lock():
                print('第{}次循环最合适的策略：'.format(self.G_EP.value), index, ' E:', TE)
                self.G_EP.value += 1

            count+=1

    def decision(self,att_points,j):

        angle=self.Choose_Route()
        if j==att_points[0] or j==att_points[1]:
            attack = [[1]]
        else:
            attack = [[0]]
        return [angle,attack]

    def Choose_Route(self):
        name = 'Go_left'
        lable = 'route_2'
        angle = [[0, 0.02]]
        return angle

    def First_step_strategy(self):  # FirstStep的策略
        strategy = []
        select = np.random.randint(0, 120)
        strategy.append([0, select])
        return strategy

class DataGenerate_2v2(mp.Process):
    def __init__(self,n_friend,n_enemy,EP,name):
        super(DataGenerate_2v2,self).__init__()
        self.N_F=n_friend
        self.N_E=n_enemy
        self.env=ENVM.ACEnv(n_friend,n_enemy)
        self.G_EP=EP
        self.name='P'+str(name)

    def run(self):
        sql = MSF11.Mysql(self.name)
        record_buffer = []
        count = 1

        while self.G_EP.value<MAX_EP:
            E,TE,index = 0.,0,0
            strategy=self.First_step_strategy()
            att_points=list(combinations(range(10,Step-20,Span*2),2))  # 基本最后20帧发射的导弹打不到目标；该处理减小了搜索空间
            len_att=len(att_points)
            for i in range(len_att**2):
                self.env.reset()
                ob = self.env.FirstStep(strategy)  # 随机初始位置
                record = (ob[0][0][0], ob[0][0][1], ob[0][0][2], ob[0][1][0], ob[0][1][1], ob[0][1][2],
                          ob[1][0][0], ob[1][0][1], ob[1][0][2], ob[1][1][0], ob[1][1][1], ob[1][1][2])

                if self.name=='P':
                    self.env.move()
                    fig = plt.figure()
                    ax = fig.add_subplot(111, projection='3d')
                    self.env.render(fig, ax)

                for j in range(Step):
                    ob = self.env.step(self.decision(att_points[i//len_att],att_points[i%len_att],j,strategy))
                    if self.name=='P':
                        self.env.render(fig, ax)
                    E += ob[0][2][0]    # 一回合的摧毁目标期望
                if self.name=='P':
                    plt.close()
                if E>TE:
                    TE=E
                    index=i
                E=0
                record+=(index,TE)

            record_buffer.append(record)
            if count%BufferSpan==0:
                sql.Insert_record2(DatabaseName,record_buffer)
                record_buffer.clear()
            with self.G_EP.get_lock():
                print('第{}次循环最合适的策略：'.format(self.G_EP.value), index, ' E:', TE)
                self.G_EP.value += 1

            count+=1

    def decision(self,att_points,att_points2,j,stra):
        angle=self.Choose_Route()
        attack=[[0,0],[0,0]]
        if j==att_points[0] or j==att_points[1]:
            attack[0][stra[0][0]]=1
        if j==att_points2[0] or j==att_points2[1]:
            attack[1][stra[1][0]] = 1

        return [angle,attack]

    def Choose_Route(self):
        name = 'Go_right'
        lable = 'route_12'
        angle = [[0, 0], [0, 0.02]]
        return angle

    def First_step_strategy(self): # FirstStep的策略
        strategy=[]
        select1 = np.random.randint(0, 120)
        select2=np.random.randint(0,120)
        select_target=np.random.randint(0,2,2)
        strategy.append([select_target[0],select1])
        strategy.append([select_target[1],select2])

        return strategy

if __name__=='__main__':
    EP = mp.Value('i', 0)
    workers = [DataGenerate_2v2(2,2,EP,i) for i in range(mp.cpu_count())]  #mp.cpu_count()
    [w.start() for w in workers]
    [w.join() for w in workers]

